{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import reader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1. 定义相关的参数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "DATA_PATH = \"../../datasets/PTB_data\"\n",
    "HIDDEN_SIZE = 200\n",
    "NUM_LAYERS = 2\n",
    "VOCAB_SIZE = 10000\n",
    "\n",
    "LEARNING_RATE = 1.0\n",
    "TRAIN_BATCH_SIZE = 20\n",
    "TRAIN_NUM_STEP = 35\n",
    "\n",
    "EVAL_BATCH_SIZE = 1\n",
    "EVAL_NUM_STEP = 1\n",
    "NUM_EPOCH = 2\n",
    "KEEP_PROB = 0.5\n",
    "MAX_GRAD_NORM = 5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2. 定义一个类来描述模型结构。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class PTBModel(object):\n",
    "    def __init__(self, is_training, batch_size, num_steps):\n",
    "        \n",
    "        self.batch_size = batch_size\n",
    "        self.num_steps = num_steps\n",
    "        \n",
    "        # 定义输入层。\n",
    "        self.input_data = tf.placeholder(tf.int32, [batch_size, num_steps])\n",
    "        self.targets = tf.placeholder(tf.int32, [batch_size, num_steps])\n",
    "        \n",
    "        # 定义使用LSTM结构及训练时使用dropout。\n",
    "        lstm_cell = tf.contrib.rnn.BasicLSTMCell(HIDDEN_SIZE)\n",
    "        if is_training:\n",
    "            lstm_cell = tf.contrib.rnn.DropoutWrapper(lstm_cell, output_keep_prob=KEEP_PROB)\n",
    "        cell = tf.contrib.rnn.MultiRNNCell([lstm_cell]*NUM_LAYERS)\n",
    "        \n",
    "        # 初始化最初的状态。\n",
    "        self.initial_state = cell.zero_state(batch_size, tf.float32)\n",
    "        embedding = tf.get_variable(\"embedding\", [VOCAB_SIZE, HIDDEN_SIZE])\n",
    "        \n",
    "        # 将原本单词ID转为单词向量。\n",
    "        inputs = tf.nn.embedding_lookup(embedding, self.input_data)\n",
    "        \n",
    "        if is_training:\n",
    "            inputs = tf.nn.dropout(inputs, KEEP_PROB)\n",
    "\n",
    "        # 定义输出列表。\n",
    "        outputs = []\n",
    "        state = self.initial_state\n",
    "        with tf.variable_scope(\"RNN\"):\n",
    "            for time_step in range(num_steps):\n",
    "                if time_step > 0: tf.get_variable_scope().reuse_variables()\n",
    "                cell_output, state = cell(inputs[:, time_step, :], state)\n",
    "                outputs.append(cell_output) \n",
    "        output = tf.reshape(tf.concat(outputs, 1), [-1, HIDDEN_SIZE])\n",
    "        weight = tf.get_variable(\"weight\", [HIDDEN_SIZE, VOCAB_SIZE])\n",
    "        bias = tf.get_variable(\"bias\", [VOCAB_SIZE])\n",
    "        logits = tf.matmul(output, weight) + bias\n",
    "        \n",
    "        # 定义交叉熵损失函数和平均损失。\n",
    "        loss = tf.contrib.legacy_seq2seq.sequence_loss_by_example(\n",
    "            [logits],\n",
    "            [tf.reshape(self.targets, [-1])],\n",
    "            [tf.ones([batch_size * num_steps], dtype=tf.float32)])\n",
    "        self.cost = tf.reduce_sum(loss) / batch_size\n",
    "        self.final_state = state\n",
    "        \n",
    "        # 只在训练模型时定义反向传播操作。\n",
    "        if not is_training: return\n",
    "        trainable_variables = tf.trainable_variables()\n",
    "\n",
    "        # 控制梯度大小，定义优化方法和训练步骤。\n",
    "        grads, _ = tf.clip_by_global_norm(tf.gradients(self.cost, trainable_variables), MAX_GRAD_NORM)\n",
    "        optimizer = tf.train.GradientDescentOptimizer(LEARNING_RATE)\n",
    "        self.train_op = optimizer.apply_gradients(zip(grads, trainable_variables))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3. 使用给定的模型model在数据data上运行train_op并返回在全部数据上的perplexity值"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def run_epoch(session, model, data, train_op, output_log, epoch_size):\n",
    "    total_costs = 0.0\n",
    "    iters = 0\n",
    "    state = session.run(model.initial_state)\n",
    "\n",
    "    # 训练一个epoch。\n",
    "    for step in range(epoch_size):\n",
    "        x, y = session.run(data)\n",
    "        cost, state, _ = session.run([model.cost, model.final_state, train_op],\n",
    "                                        {model.input_data: x, model.targets: y, model.initial_state: state})\n",
    "        total_costs += cost\n",
    "        iters += model.num_steps\n",
    "\n",
    "        if output_log and step % 100 == 0:\n",
    "            print(\"After %d steps, perplexity is %.3f\" % (step, np.exp(total_costs / iters)))\n",
    "    return np.exp(total_costs / iters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4. 定义主函数并执行。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "In iteration: 1\n",
      "After 0 steps, perplexity is 9979.425\n",
      "After 100 steps, perplexity is 1419.514\n",
      "After 200 steps, perplexity is 1069.710\n",
      "After 300 steps, perplexity is 890.482\n",
      "After 400 steps, perplexity is 779.052\n",
      "After 500 steps, perplexity is 702.706\n",
      "After 600 steps, perplexity is 645.688\n",
      "After 700 steps, perplexity is 595.862\n",
      "After 800 steps, perplexity is 551.698\n",
      "After 900 steps, perplexity is 517.378\n",
      "After 1000 steps, perplexity is 490.698\n",
      "After 1100 steps, perplexity is 465.219\n",
      "After 1200 steps, perplexity is 444.531\n",
      "After 1300 steps, perplexity is 425.488\n",
      "Epoch: 1 Validation Perplexity: 235.968\n",
      "In iteration: 2\n",
      "After 0 steps, perplexity is 353.259\n",
      "After 100 steps, perplexity is 244.039\n",
      "After 200 steps, perplexity is 249.303\n",
      "After 300 steps, perplexity is 249.369\n",
      "After 400 steps, perplexity is 246.174\n",
      "After 500 steps, perplexity is 243.266\n",
      "After 600 steps, perplexity is 242.566\n",
      "After 700 steps, perplexity is 239.728\n",
      "After 800 steps, perplexity is 235.205\n",
      "After 900 steps, perplexity is 232.611\n",
      "After 1000 steps, perplexity is 231.105\n",
      "After 1100 steps, perplexity is 227.587\n",
      "After 1200 steps, perplexity is 225.045\n",
      "After 1300 steps, perplexity is 222.219\n",
      "Epoch: 2 Validation Perplexity: 181.089\n",
      "Test Perplexity: 176.798\n"
     ]
    }
   ],
   "source": [
    "def main():\n",
    "    train_data, valid_data, test_data, _ = reader.ptb_raw_data(DATA_PATH)\n",
    "\n",
    "    # 计算一个epoch需要训练的次数\n",
    "    train_data_len = len(train_data)\n",
    "    train_batch_len = train_data_len // TRAIN_BATCH_SIZE\n",
    "    train_epoch_size = (train_batch_len - 1) // TRAIN_NUM_STEP\n",
    "\n",
    "    valid_data_len = len(valid_data)\n",
    "    valid_batch_len = valid_data_len // EVAL_BATCH_SIZE\n",
    "    valid_epoch_size = (valid_batch_len - 1) // EVAL_NUM_STEP\n",
    "\n",
    "    test_data_len = len(test_data)\n",
    "    test_batch_len = test_data_len // EVAL_BATCH_SIZE\n",
    "    test_epoch_size = (test_batch_len - 1) // EVAL_NUM_STEP\n",
    "\n",
    "    initializer = tf.random_uniform_initializer(-0.05, 0.05)\n",
    "    with tf.variable_scope(\"language_model\", reuse=None, initializer=initializer):\n",
    "        train_model = PTBModel(True, TRAIN_BATCH_SIZE, TRAIN_NUM_STEP)\n",
    "\n",
    "    with tf.variable_scope(\"language_model\", reuse=True, initializer=initializer):\n",
    "        eval_model = PTBModel(False, EVAL_BATCH_SIZE, EVAL_NUM_STEP)\n",
    "\n",
    "    # 训练模型。\n",
    "    with tf.Session() as session:\n",
    "        tf.global_variables_initializer().run()\n",
    "\n",
    "        train_queue = reader.ptb_producer(train_data, train_model.batch_size, train_model.num_steps)\n",
    "        eval_queue = reader.ptb_producer(valid_data, eval_model.batch_size, eval_model.num_steps)\n",
    "        test_queue = reader.ptb_producer(test_data, eval_model.batch_size, eval_model.num_steps)\n",
    "\n",
    "        coord = tf.train.Coordinator()\n",
    "        threads = tf.train.start_queue_runners(sess=session, coord=coord)\n",
    "\n",
    "        for i in range(NUM_EPOCH):\n",
    "            print(\"In iteration: %d\" % (i + 1))\n",
    "            run_epoch(session, train_model, train_queue, train_model.train_op, True, train_epoch_size)\n",
    "\n",
    "            valid_perplexity = run_epoch(session, eval_model, eval_queue, tf.no_op(), False, valid_epoch_size)\n",
    "            print(\"Epoch: %d Validation Perplexity: %.3f\" % (i + 1, valid_perplexity))\n",
    "\n",
    "        test_perplexity = run_epoch(session, eval_model, test_queue, tf.no_op(), False, test_epoch_size)\n",
    "        print(\"Test Perplexity: %.3f\" % test_perplexity)\n",
    "\n",
    "        coord.request_stop()\n",
    "        coord.join(threads)\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
